#!/usr/bin/env python3

import rospy
import time
import math
import numpy as np
import os  # For path handling

# ROS Messages
from std_msgs.msg import Int32, Float64, Header
from geometry_msgs.msg import PoseStamped, TwistStamped, Vector3Stamped
from mavros_msgs.msg import PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode

# TF Transformations for Euler angles
import tf.transformations as tf

# YOLO Detection Message (Assumed)
# Ensure this matches your YOLO output message type
from vision_msgs.msg import Detection2DArray

# Import calibration utility
from calibrate_parser import initialize_calibration


class UAVIBVSController:
    def __init__(self):
        rospy.init_node('uav_ibvs_controller_node', anonymous=True)

        # --- 1. Parameters Loading ---
        self.uav_id = rospy.get_param('~uav_id', 1)  # This node assumes a single UAV, but kept for consistency
        rospy.loginfo(f"Initializing IBVS controller for UAV {self.uav_id}")

        self.control_frequency = rospy.get_param('~control_frequency', 20.0)  # Hz
        self.rate = rospy.Rate(self.control_frequency)

        # Camera calibration file path
        script_dir = rospy.get_param('~script_dir')  # Must be set in launch file
        calibration_file_path = os.path.join(script_dir, 'calibrate.txt')
        self.K_matrix, self.image_width, self.image_height = initialize_calibration(calibration_file_path)
        if self.K_matrix is None or self.image_width == 0 or self.image_height == 0:
            rospy.logerr("Failed to load camera calibration! Shutting down.")
            rospy.signal_shutdown("Camera calibration error.")
            return

        self.image_center_x = self.image_width / 2.0
        self.image_center_y = self.image_height / 2.0

        # IBVS Control Gains (NEED CAREFUL TUNING!)
        self.desired_target_area_pixels = rospy.get_param('~desired_target_area_pixels',
                                                          150 * 150)  # e.g., 150x150 pixels

        # These gains are proportional (P) controllers based on image errors
        self.k_yaw_ibvs = rospy.get_param('~k_yaw_ibvs', 0.005)  # Yaw angular velocity gain (for X-axis centering)
        self.k_pitch_gimbal_ibvs = rospy.get_param('~k_pitch_gimbal_ibvs',
                                                   0.05)  # Gimbal pitch angle gain (for Y-axis centering)
        self.k_forward_speed_ibvs = rospy.get_param('~k_forward_speed_ibvs',
                                                    0.0005)  # Forward/backward speed gain (for distance control)

        # Speed and Angle Limits
        self.max_yaw_rate = rospy.get_param('~max_yaw_rate', 0.5)  # rad/s
        self.max_forward_speed = rospy.get_param('~max_forward_speed', 3.0)  # m/s (max forward)
        self.max_backward_speed = rospy.get_param('~max_backward_speed', 1.0)  # m/s (max backward)
        self.max_gimbal_pitch_angle = rospy.get_param('~max_gimbal_pitch_angle', 30.0)  # deg
        self.min_gimbal_pitch_angle = rospy.get_param('~min_gimbal_pitch_angle', -90.0)  # deg
        self.desired_vertical_speed = rospy.get_param('~desired_vertical_speed',
                                                      0.0)  # m/s (set to 0 for altitude hold)

        self.target_lost_timeout = rospy.get_param('~target_lost_timeout', 1.0)  # seconds

        # --- 2. State Variables ---
        self.uav_current_pose = np.zeros(6)  # [x, y, z, roll, pitch, yaw] (radians for angles)
        self.px4_state = State()  # MAVROS state message
        self.target_bbox_yolo = None  # [x_min, y_min, x_max, y_max] from YOLO
        self.last_detection_time = rospy.Time(0)
        self.target_is_tracked = False
        self.current_gimbal_pitch_deg = 0.0  # From gimbal feedback

        self.current_flight_state = "INIT"  # Custom state machine for drone control
        self.auto_arm = True  # Whether to auto-arm the drone
        self.takeoff_altitude = rospy.get_param('~takeoff_altitude', 1.0)  # meters
        self.initial_takeoff_pos = None  # Stored position at takeoff command

        # --- 3. ROS Publishers & Subscribers ---
        # MAVROS Publisher for velocity/yawrate commands (PositionTarget)
        self.mavros_setpoint_pub = rospy.Publisher(f'/uav{self.uav_id}/mavros/setpoint_raw/local', PositionTarget,
                                                   queue_size=1)

        # Gimbal control publisher
        # Assumed gimbal API takes Vector3Stamped (x: pitch, y: roll, z: yaw) in degrees
        self.gimbal_cmd_pub = rospy.Publisher('/gimbal_controller/command/euler', Vector3Stamped, queue_size=1)

        # Subscribers
        rospy.Subscriber(f'/uav{self.uav_id}/mavros/local_position/pose', PoseStamped, self._uav_pose_cb)
        rospy.Subscriber(f'/uav{self.uav_id}/mavros/state', State, self._px4_state_cb)
        rospy.Subscriber('/yolo/detections', Detection2DArray, self._yolo_detections_cb)
        rospy.Subscriber('/gimbal_controller/pitch_angle_feedback', Float64,
                         self._gimbal_pitch_feedback_cb)  # Assuming feedback
        rospy.Subscriber('/command', Int32, self._command_cb)  # Ground station command

        # MAVROS Service Proxies
        rospy.wait_for_service(f'/uav{self.uav_id}/mavros/cmd/arming')
        self.arming_client = rospy.ServiceProxy(f'/uav{self.uav_id}/mavros/cmd/arming', CommandBool)
        rospy.wait_for_service(f'/uav{self.uav_id}/mavros/set_mode')
        self.set_mode_client = rospy.ServiceProxy(f'/uav{self.uav_id}/mavros/set_mode', SetMode)

        rospy.loginfo("UAV IBVS Controller Initialized.")

    # --- ROS Callback Functions ---
    def _uav_pose_cb(self, msg: PoseStamped):
        self.uav_current_pose[0] = msg.pose.position.x
        self.uav_current_pose[1] = msg.pose.position.y
        self.uav_current_pose[2] = msg.pose.position.z

        # Convert quaternion to Euler angles (roll, pitch, yaw)
        q = [msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z, msg.pose.orientation.w]
        (self.uav_current_pose[3], self.uav_current_pose[4], self.uav_current_pose[5]) = tf.euler_from_quaternion(q)

    def _px4_state_cb(self, msg: State):
        self.px4_state = msg

    def _yolo_detections_cb(self, msg: Detection2DArray):
        if msg.detections:
            # Take the largest detection as the target
            largest_detection = max(msg.detections, key=lambda d: d.bbox.size_x * d.bbox.size_y)
            bbox = largest_detection.bbox
            self.target_bbox_yolo = [bbox.center.x - bbox.size_x / 2,  # x_min
                                     bbox.center.y - bbox.size_y / 2,  # y_min
                                     bbox.center.x + bbox.size_x / 2,  # x_max
                                     bbox.center.y + bbox.size_y / 2]  # y_max
            self.target_is_tracked = True
            self.last_detection_time = rospy.Time.now()
        else:
            self.target_is_tracked = False

    def _gimbal_pitch_feedback_cb(self, msg: Float64):
        self.current_gimbal_pitch_deg = msg.data  # Assuming feedback is in degrees

    def _command_cb(self, msg: Int32):
        cmd = msg.data
        if cmd == 0:
            self.current_flight_state = 'INIT'
        elif cmd == 1:
            self.current_flight_state = 'ARM'
        elif cmd == 2:
            self.current_flight_state = 'TAKEOFF'
        elif cmd == 3:
            self.current_flight_state = 'LAND'
        elif cmd == 4:
            self.current_flight_state = 'RESET_POS'  # Go to initial hover position (e.g., origin)
        elif cmd == 5:
            self.current_flight_state = 'START_IBVS_TRACKING'
        elif cmd == 7:
            self.current_flight_state = 'WAITING'
        else:
            rospy.logwarn(f"Unknown command: {cmd}. Maintaining current state: {self.current_flight_state}")
        rospy.loginfo(f"Received command: {cmd}, transitioning to state: {self.current_flight_state}")

    # --- MAVROS Service Calls ---
    def _arm_drone(self) -> bool:
        if self.arming_client(True):
            rospy.loginfo(f"UAV {self.uav_id} armed successfully!")
            return True
        else:
            rospy.logerr(f"UAV {self.uav_id} arming failed!")
            return False

    def _disarm_drone(self) -> bool:
        if self.arming_client(False):
            rospy.loginfo(f"UAV {self.uav_id} disarmed successfully!")
            return True
        else:
            rospy.logerr(f"UAV {self.uav_id} disarming failed!")
            return False

    def _set_offboard_mode(self) -> bool:
        if self.set_mode_client(0, "OFFBOARD").mode_sent:
            rospy.loginfo(f"UAV {self.uav_id} set to OFFBOARD mode.")
            return True
        else:
            rospy.logerr(f"UAV {self.uav_id} failed to set OFFBOARD mode!")
            return False

    # --- MAVROS Command Publishing Functions ---
    def _publish_velocity_yawrate_cmd(self, vx: float, vy: float, vz: float, yaw_rate: float):
        """
        Publishes velocity and yaw rate commands to MAVROS PositionTarget.
        vx, vy, vz: in local NED frame (PX4 default)
        yaw_rate: in rad/s
        """
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = PositionTarget.FRAME_LOCAL_NED  # Local NED frame

        # Type mask to control velocity and yaw_rate only
        msg.type_mask = PositionTarget.IGNORE_PX + \
                        PositionTarget.IGNORE_PY + \
                        PositionTarget.IGNORE_PZ + \
                        PositionTarget.IGNORE_AFX + \
                        PositionTarget.IGNORE_AFY + \
                        PositionTarget.IGNORE_AFZ + \
                        PositionTarget.IGNORE_YAW  # Ignore yaw position, control yaw rate

        msg.velocity.x = vx
        msg.velocity.y = vy
        msg.velocity.z = vz
        msg.yaw_rate = yaw_rate
        self.mavros_setpoint_pub.publish(msg)

    def _publish_position_cmd(self, x: float, y: float, z: float):
        """
        Publishes position commands to MAVROS PositionTarget.
        x, y, z: in local NED frame (PX4 default)
        """
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = PositionTarget.FRAME_LOCAL_NED

        # Type mask to control position only
        msg.type_mask = PositionTarget.IGNORE_VX + \
                        PositionTarget.IGNORE_VY + \
                        PositionTarget.IGNORE_VZ + \
                        PositionTarget.IGNORE_AFX + \
                        PositionTarget.IGNORE_AFY + \
                        PositionTarget.IGNORE_AFZ + \
                        PositionTarget.IGNORE_YAW_AND_RATE

        msg.position.x = x
        msg.position.y = y
        msg.position.z = z
        self.mavros_setpoint_pub.publish(msg)

    def _publish_posxy_velz_cmd(self, x: float, y: float, vz: float):
        """
        Publishes position X, Y and velocity Z commands.
        Used for landing.
        """
        msg = PositionTarget()
        msg.header.stamp = rospy.Time.now()
        msg.coordinate_frame = PositionTarget.FRAME_LOCAL_NED

        msg.type_mask = PositionTarget.IGNORE_VX + \
                        PositionTarget.IGNORE_VY + \
                        PositionTarget.IGNORE_PZ + \
                        PositionTarget.IGNORE_AFX + \
                        PositionTarget.IGNORE_AFY + \
                        PositionTarget.IGNORE_AFZ + \
                        PositionTarget.IGNORE_YAW_AND_RATE

        msg.position.x = x
        msg.position.y = y
        msg.velocity.z = vz
        self.mavros_setpoint_pub.publish(msg)

    def _send_gimbal_command(self, pitch_angle_deg: float, roll_angle_deg: float):
        """
        Publishes desired gimbal Euler angles (pitch and roll) to the gimbal controller.
        Angles are in degrees.
        """
        gimbal_cmd = Vector3Stamped()
        gimbal_cmd.header.stamp = rospy.Time.now()
        gimbal_cmd.header.frame_id = "gimbal_link"  # Or your gimbal's specific frame
        gimbal_cmd.vector.x = pitch_angle_deg  # Common convention: X for Pitch
        gimbal_cmd.vector.y = roll_angle_deg  # Common convention: Y for Roll
        gimbal_cmd.vector.z = 0.0  # Z for Yaw, typically not controlled by IBVS for gimbal
        self.gimbal_cmd_pub.publish(gimbal_cmd)

    # --- Auxiliary Functions ---
    def _distance_3d(self, pos1: np.ndarray, pos2: np.ndarray) -> float:
        """Calculates Euclidean distance between two 3D points."""
        return np.sqrt((pos1[0] - pos2[0]) ** 2 + (pos1[1] - pos2[1]) ** 2 + (pos1[2] - pos2[2]) ** 2)

    # --- IBVS Control Logic ---
    def _run_ibvs_control(self):
        # 1. Check target status
        if not self.target_is_tracked and (
                rospy.Time.now() - self.last_detection_time).to_sec() > self.target_lost_timeout:
            rospy.logwarn("Target lost for too long. Hovering drone and centering gimbal.")
            self._publish_velocity_yawrate_cmd(0.0, 0.0, self.desired_vertical_speed, 0.0)
            self._send_gimbal_command(0.0, 0.0)  # Reset gimbal to horizontal
            # Optionally, switch state to 'SEARCH' or 'WAITING'
            self.current_flight_state = 'WAITING'
            return

        if self.target_bbox_yolo is None:  # No detection ever
            rospy.loginfo_throttle(1.0, "Waiting for initial YOLO detection to start IBVS tracking.")
            self._publish_velocity_yawrate_cmd(0.0, 0.0, self.desired_vertical_speed, 0.0)  # Hover
            self._send_gimbal_command(0.0, 0.0)
            return

        # --- Delayed Kalman Filter (DKF) Placeholder ---
        # IMPORTANT: This is the most crucial part for robust performance in the paper!
        # Currently, it directly uses the raw YOLO bbox, which introduces latency and noise.
        # For a real implementation, you would replace this with a filter
        # that integrates UAV IMU data with YOLO detections to provide
        # a smooth, low-latency estimate of the target's image coordinates and area.
        # Example (conceptual):
        # estimated_x, estimated_y, estimated_area = self.dkf.update(self.target_bbox_yolo, self.uav_current_pose, rospy.Time.now())

        # For now, directly use current YOLO data (will be noisy and delayed)
        current_x_pixel = (self.target_bbox_yolo[0] + self.target_bbox_yolo[2]) / 2.0
        current_y_pixel = (self.target_bbox_yolo[1] + self.target_bbox_yolo[3]) / 2.0
        current_area_pixels = (self.target_bbox_yolo[2] - self.target_bbox_yolo[0]) * \
                              (self.target_bbox_yolo[3] - self.target_bbox_yolo[1])
        # --- END DKF Placeholder ---

        # 2. Calculate Image Errors
        # Error for X-axis centering (controls Yaw)
        error_x_pixel = current_x_pixel - self.image_center_x
        # Positive error means target is to the right of center.

        # Error for Y-axis centering (controls Gimbal Pitch)
        error_y_pixel = current_y_pixel - self.image_center_y
        # Positive error means target is below center.

        # Error for distance control (controls Forward/Backward speed)
        error_area = self.desired_target_area_pixels - current_area_pixels
        # Positive error means target is too small (needs to move forward).

        # 3. Calculate Control Commands

        # UAV Yaw Rate Command (based on X-axis error)
        # If target is right (error_x_pixel > 0), need to yaw right (positive yaw_rate)
        yaw_rate_cmd = self.k_yaw_ibvs * error_x_pixel
        yaw_rate_cmd = np.clip(yaw_rate_cmd, -self.max_yaw_rate, self.max_yaw_rate)

        # Gimbal Pitch Angle Command (based on Y-axis error)
        # If target is down (error_y_pixel > 0), need to pitch gimbal down (more negative angle)
        # This is a P-controller aiming for a desired pitch angle
        # Note: 'desired_gimbal_pitch' is the *absolute* angle in degrees
        desired_gimbal_pitch = self.k_pitch_gimbal_ibvs * error_y_pixel
        desired_gimbal_pitch = np.clip(desired_gimbal_pitch, self.min_gimbal_pitch_angle, self.max_gimbal_pitch_angle)

        # UAV Forward/Backward Speed Command (based on Area error)
        # If target is too small (error_area > 0), need to move forward (positive vx)
        forward_speed_cmd = self.k_forward_speed_ibvs * error_area
        forward_speed_cmd = np.clip(forward_speed_cmd, -self.max_backward_speed, self.max_forward_speed)

        # Gimbal Roll Compensation (to keep image horizontal)
        # UAV's current roll angle (self.uav_current_pose[3]) is in radians
        # Gimbal roll command should be the negative of drone's roll to counteract it
        gimbal_roll_cmd = -np.degrees(self.uav_current_pose[3])

        # 4. Publish Commands
        # For MAVROS PositionTarget, x is forward, y is right, z is down in local NED.
        # Our forward_speed_cmd is along the drone's body x-axis.
        # Our desired_vertical_speed is along drone's body z-axis.
        # Our yaw_rate_cmd is along drone's body z-axis.
        self._publish_velocity_yawrate_cmd(forward_speed_cmd, 0.0, self.desired_vertical_speed, yaw_rate_cmd)
        self._send_gimbal_command(desired_gimbal_pitch, gimbal_roll_cmd)

        rospy.loginfo_throttle(0.5,
                               f"Errors: X={error_x_pixel:.1f}, Y={error_y_pixel:.1f}, Area={error_area:.1f} | " +
                               f"Cmds: Vx={forward_speed_cmd:.2f}, YawRate={yaw_rate_cmd:.2f}, " +
                               f"Gimbal Pitch={desired_gimbal_pitch:.1f} deg, Gimbal Roll={gimbal_roll_cmd:.1f} deg"
                               )

    # --- Main Control Loop / State Machine ---
    def run(self):
        rospy.loginfo("UAV IBVS Controller started. Waiting for commands.")

        # Ensure OFFBOARD mode is set before arming/takeoff
        # Send initial setpoints to prevent timeout
        for _ in range(100):
            self._publish_velocity_yawrate_cmd(0.0, 0.0, 0.0, 0.0)
            self.rate.sleep()

        last_mode_request = rospy.Time.now()
        while not rospy.is_shutdown():
            # Attempt to set OFFBOARD mode if not already in it
            if self.px4_state.mode != "OFFBOARD" and (rospy.Time.now() - last_mode_request).to_sec() > 1.0:
                if self._set_offboard_mode():
                    rospy.loginfo("Successfully requested OFFBOARD mode.")
                last_mode_request = rospy.Time.now()

            # State Machine Logic
            if self.current_flight_state == "INIT":
                self._publish_velocity_yawrate_cmd(0.0, 0.0, 0.0, 0.0)  # Hover / stay put
                self._send_gimbal_command(0.0, 0.0)  # Center gimbal
                # Wait for ARM command

            elif self.current_flight_state == "ARM":
                if not self.px4_state.armed and self.auto_arm:
                    if self._arm_drone():
                        self.current_flight_state = "WAITING"  # Armed, ready for takeoff

            elif self.current_flight_state == "TAKEOFF":
                if self.initial_takeoff_pos is None:
                    # Store current XY and target Z for takeoff
                    self.initial_takeoff_pos = [self.uav_current_pose[0], self.uav_current_pose[1],
                                                self.takeoff_altitude]
                    rospy.loginfo(
                        f"Taking off to {self.takeoff_altitude}m from {[self.uav_current_pose[0]:.2f}, {self.uav_current_pose[1]:.2f]}")

                # Command position
                self._publish_position_cmd(self.initial_takeoff_pos[0], self.initial_takeoff_pos[1],
                                           self.initial_takeoff_pos[2])

                # Check if reached target altitude
                if abs(self.uav_current_pose[2] - self.initial_takeoff_pos[2]) < 0.1:  # 10cm tolerance
                    rospy.loginfo("Takeoff complete! Hovering.")
                    self.current_flight_state = "WAITING"
                    self.initial_takeoff_pos = None  # Reset for next takeoff if needed

            elif self.current_flight_state == "RESET_POS":
                target_reset_pos = [0.0, 0.0, self.takeoff_altitude]  # Go to origin at takeoff altitude
                self._publish_position_cmd(target_reset_pos[0], target_reset_pos[1], target_reset_pos[2])

                if self._distance_3d(self.uav_current_pose[:3], np.array(target_reset_pos)) < 0.2:  # 20cm tolerance
                    rospy.loginfo("Reset position reached.")
                    self.current_flight_state = "WAITING"

            elif self.current_flight_state == "START_IBVS_TRACKING":
                self._run_ibvs_control()  # Execute IBVS logic

            elif self.current_flight_state == "LAND":
                # First, ensure we're directly above the landing spot (current X, Y)
                land_target_xy = self.uav_current_pose[0:2]  # Land at current XY position
                rospy.loginfo(f"Initiating landing at current XY: {land_target_xy[0]:.2f}, {land_target_xy[1]:.2f}")
                self.current_flight_state = "LANDING_DESCEND"

            elif self.current_flight_state == "LANDING_DESCEND":
                land_target_xy = self.uav_current_pose[0:2]  # Keep landing at current XY

                # Descend slowly
                if self.uav_current_pose[2] < 0.2:  # Near ground (adjust threshold as needed for your drone's height)
                    # When close to ground, command a slightly faster descent to touch down
                    self._publish_posxy_velz_cmd(land_target_xy[0], land_target_xy[1],
                                                 0.8)  # vz=0.8 (downward, since NED is Z-down)
                else:
                    # Otherwise, descend steadily
                    self._publish_posxy_velz_cmd(land_target_xy[0], land_target_xy[1], 0.2)  # vz=0.2 (downward)

                if self.uav_current_pose[2] > -0.1 and self.uav_current_pose[2] < 0.1:  # Near Z=0 (on ground)
                    rospy.loginfo("Landing complete. Disarming.")
                    self.current_flight_state = "DISARM"
                    # Reset gimbal
                    self._send_gimbal_command(0.0, 0.0)

            elif self.current_flight_state == "DISARM":
                if self._disarm_drone():
                    rospy.loginfo("UAV disarmed. Node shutting down.")
                    rospy.signal_shutdown("Disarmed.")

            elif self.current_flight_state == "WAITING":
                self._publish_velocity_yawrate_cmd(0.0, 0.0, 0.0, 0.0)  # Hover
                self._send_gimbal_command(0.0, 0.0)  # Center gimbal

            self.rate.sleep()  # Maintain control loop frequency


if __name__ == "__main__":
    try:
        controller = UAVIBVSController()
        controller.run()
    except rospy.ROSInterruptException:
        pass
    except Exception as e:
        rospy.logerr(f"An unhandled error occurred: {e}")
